"""
FAISS向量索引管理工具
提供高效的人脸特征向量检索功能
"""
import os
import pickle
from typing import List, Tuple, Optional
import numpy as np
import faiss
from utils.logger import app_logger
from config.settings import settings


class FaissIndexManager:
    """FAISS索引管理器（单例模式）"""

    _instance = None

    def __new__(cls):
        if cls._instance is None:
            cls._instance = super().__new__(cls)
            cls._instance._initialized = False
        return cls._instance

    def __init__(self):
        if self._initialized:
            return

        self.dimension = settings.faiss_dimension
        self.index_path = settings.faiss_index_path
        self.face_id_map_path = self.index_path.replace('.bin', '_map.pkl')

        # 创建索引：使用内积（归一化向量的内积等于余弦相似度）
        self.index: Optional[faiss.IndexFlatIP] = None
        self.face_id_map: List[str] = []  # 索引位置 -> face_id 映射

        self._load_or_create_index()
        self._initialized = True

    def _load_or_create_index(self):
        """加载或创建索引"""
        if os.path.exists(self.index_path) and os.path.exists(self.face_id_map_path):
            try:
                # 加载索引
                self.index = faiss.read_index(self.index_path)
                # 加载映射
                with open(self.face_id_map_path, 'rb') as f:
                    self.face_id_map = pickle.load(f)
                app_logger.info(f"FAISS索引加载成功: {self.index.ntotal}条记录")
            except Exception as e:
                app_logger.error(f"加载FAISS索引失败: {e}", exc_info=True)
                self._create_new_index()
        else:
            self._create_new_index()

    def _create_new_index(self):
        """创建新索引"""
        self.index = faiss.IndexFlatIP(self.dimension)  # 内积索引
        self.face_id_map = []
        app_logger.info("创建新的FAISS索引")

    def add_vector(self, face_id: str, feature_vector: np.ndarray) -> bool:
        """
        添加单个向量

        Args:
            face_id: 人脸ID
            feature_vector: 特征向量(512维)

        Returns:
            是否成功
        """
        try:
            # 归一化
            vector = self._normalize_vector(feature_vector)

            # 添加到索引
            self.index.add(vector.reshape(1, -1).astype('float32'))
            self.face_id_map.append(face_id)

            app_logger.debug(f"向量添加成功: face_id={face_id}")
            return True

        except Exception as e:
            app_logger.error(f"添加向量失败: {e}", exc_info=True)
            return False

    def batch_add_vectors(self, face_ids: List[str], feature_vectors: np.ndarray) -> int:
        """
        批量添加向量

        Args:
            face_ids: 人脸ID列表
            feature_vectors: 特征向量数组 (n, 512)

        Returns:
            成功添加的数量
        """
        try:
            if len(face_ids) != feature_vectors.shape[0]:
                raise ValueError("face_ids数量与特征向量数量不匹配")

            # 归一化
            vectors = np.array([self._normalize_vector(v) for v in feature_vectors])

            # 批量添加
            self.index.add(vectors.astype('float32'))
            self.face_id_map.extend(face_ids)

            count = len(face_ids)
            app_logger.info(f"批量添加向量成功: {count}条")
            return count

        except Exception as e:
            app_logger.error(f"批量添加向量失败: {e}", exc_info=True)
            return 0

    def search(self, query_vector: np.ndarray, top_k: int = 5) -> List[Tuple[str, float]]:
        """
        搜索最相似的向量

        Args:
            query_vector: 查询向量(512维)
            top_k: 返回前K个结果

        Returns:
            [(face_id, similarity), ...] 按相似度降序排列
        """
        try:
            if self.index.ntotal == 0:
                app_logger.warning("索引为空，无法搜索")
                return []
            
            # 验证索引完整性
            if not self.validate_index():
                app_logger.error("索引不一致，尝试从数据库重建")
                return []

            # 归一化
            vector = self._normalize_vector(query_vector)

            # 搜索
            k = min(top_k, self.index.ntotal)
            distances, indices = self.index.search(
                vector.reshape(1, -1).astype('float32'), k
            )

            # 构建结果
            results = []
            for i, idx in enumerate(indices[0]):
                if idx != -1 and idx < len(self.face_id_map):  # 增加边界检查
                    face_id = self.face_id_map[idx]
                    similarity = float(distances[0][i])  # 归一化后的内积就是余弦相似度
                    results.append((face_id, similarity))

            return results

        except Exception as e:
            app_logger.error(f"向量搜索失败: {e}", exc_info=True)
            return []

    def remove_by_face_id(self, face_id: str) -> bool:
        """
        删除指定face_id的向量（通过重建索引）

        Args:
            face_id: 人脸ID

        Returns:
            是否成功
        """
        try:
            if face_id not in self.face_id_map:
                app_logger.warning(f"face_id不存在于索引中: {face_id}")
                return False

            # 找到索引位置
            idx = self.face_id_map.index(face_id)

            # 保存旧的映射（在创建新索引前）
            old_face_id_map = self.face_id_map.copy()

            # 重建索引（FAISS不支持直接删除，需要重建）
            all_vectors = []
            new_face_ids = []
            for i in range(self.index.ntotal):
                if i != idx:
                    vector = self.index.reconstruct(int(i))
                    all_vectors.append(vector)
                    new_face_ids.append(old_face_id_map[i])

            # 创建新索引
            self._create_new_index()

            # 重新添加向量
            if all_vectors:
                vectors = np.array(all_vectors)
                self.index.add(vectors.astype('float32'))
                self.face_id_map = new_face_ids

            app_logger.info(f"删除向量成功: face_id={face_id}, 剩余{len(self.face_id_map)}条记录")
            return True

        except Exception as e:
            app_logger.error(f"删除向量失败: {e}", exc_info=True)
            return False

    def save(self) -> bool:
        """保存索引到磁盘"""
        try:
            # 验证索引完整性
            if not self.validate_index():
                app_logger.error("索引不一致，拒绝保存")
                return False
            
            # 确保目录存在
            os.makedirs(os.path.dirname(self.index_path), exist_ok=True)

            # 保存索引
            faiss.write_index(self.index, self.index_path)

            # 保存映射
            with open(self.face_id_map_path, 'wb') as f:
                pickle.dump(self.face_id_map, f)

            app_logger.info(f"FAISS索引保存成功: {self.index.ntotal}条记录")
            return True

        except Exception as e:
            app_logger.error(f"保存FAISS索引失败: {e}", exc_info=True)
            return False

    def get_total_count(self) -> int:
        """获取索引中的向量数量"""
        return self.index.ntotal if self.index else 0

    def validate_index(self) -> bool:
        """
        验证索引完整性
        确保索引向量数量与映射数量一致
        
        Returns:
            是否有效
        """
        try:
            if self.index is None:
                app_logger.warning("索引未初始化")
                return False
            
            index_count = self.index.ntotal
            map_count = len(self.face_id_map)
            
            if index_count != map_count:
                app_logger.error(
                    f"索引不一致！向量数量={index_count}, 映射数量={map_count}"
                )
                return False
            
            app_logger.debug(f"索引验证通过: {index_count}条记录")
            return True
            
        except Exception as e:
            app_logger.error(f"索引验证失败: {e}", exc_info=True)
            return False

    @staticmethod
    def _normalize_vector(vector: np.ndarray) -> np.ndarray:
        """L2归一化向量"""
        norm = np.linalg.norm(vector)
        if norm == 0:
            return vector
        return vector / norm

    def rebuild_from_database(self, face_data: List[Tuple[str, np.ndarray]]):
        """
        从数据库数据重建索引

        Args:
            face_data: [(face_id, feature_vector), ...]
        """
        try:
            app_logger.info("开始重建FAISS索引...")

            # 创建新索引
            self._create_new_index()

            if face_data:
                face_ids = [item[0] for item in face_data]
                vectors = np.array([item[1] for item in face_data])
                self.batch_add_vectors(face_ids, vectors)

            # 保存
            self.save()

            app_logger.info(f"FAISS索引重建完成: {self.index.ntotal}条记录")

        except Exception as e:
            app_logger.error(f"重建索引失败: {e}", exc_info=True)
            raise


# 全局单例实例
faiss_manager = FaissIndexManager()
